import sys
import os
import pickle
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.figure import Figure
import matplotlib
from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout,
                             QLabel, QPushButton, QComboBox, QFileDialog, QTabWidget,
                             QTableWidget, QTableWidgetItem, QSplitter, QScrollArea,
                             QGridLayout, QGroupBox, QCheckBox, QSlider, QSpinBox)
from PyQt5.QtGui import QPixmap, QImage
from PyQt5.QtCore import Qt, QThread, pyqtSignal, QSize
import torch
from mmengine.config import Config
from mmdet.registry import DATASETS
from mmdet.evaluation import eval_map
from mmdet.apis import init_detector, inference_detector
from PIL import Image, ImageDraw, ImageFont
import cv2

# 设置matplotlib使用中文字体
matplotlib.rcParams['font.sans-serif'] = ['SimHei']
matplotlib.rcParams['axes.unicode_minus'] = False


class MatplotlibCanvas(FigureCanvas):
    def __init__(self, parent=None, width=5, height=4, dpi=100):
        self.fig = Figure(figsize=(width, height), dpi=dpi)
        self.axes = self.fig.add_subplot(111)
        super(MatplotlibCanvas, self).__init__(self.fig)


class EvaluationThread(QThread):
    finished = pyqtSignal(dict)
    progress = pyqtSignal(str)

    def __init__(self, result_path, config_path):
        super().__init__()
        self.result_path = result_path
        self.config_path = config_path

    def run(self):
        try:
            # 加载测试结果和数据集
            self.progress.emit("正在加载结果文件...")
            with open(self.result_path, 'rb') as f:
                results = pickle.load(f)

            self.progress.emit("正在加载配置文件...")
            cfg = Config.fromfile(self.config_path)

            # 构建数据集
            self.progress.emit("正在构建数据集...")
            dataset = DATASETS.build(cfg.test_dataloader.dataset)

            # 获取类别名称
            class_names = dataset.metainfo['classes']
            num_classes = len(class_names)

            self.progress.emit(f"数据集包含 {num_classes} 个类别")

            # 提取标注信息
            self.progress.emit("正在提取标注信息...")
            annotations = []
            for i in range(len(dataset)):
                data_info = dataset.get_data_info(i)
                instances = data_info.get('instances', [])

                bboxes = []
                labels = []

                for instance in instances:
                    if instance.get('ignore_flag', 0) == 0:
                        bboxes.append(instance['bbox'])
                        labels.append(instance['bbox_label'])

                ann = {
                    'bboxes': np.array(bboxes, dtype=np.float32),
                    'labels': np.array(labels, dtype=np.int64)
                }
                annotations.append(ann)

            # 转换结果格式
            if isinstance(results[0], dict) and 'pred_instances' in results[0]:
                self.progress.emit("正在转换结果格式...")
                converted_results = []
                for result in results:
                    pred_instances = result['pred_instances']
                    scores = pred_instances.get('scores', []).cpu().numpy()
                    bboxes = pred_instances.get('bboxes', []).cpu().numpy()
                    labels = pred_instances.get('labels', []).cpu().numpy()

                    result_list = [np.zeros((0, 5), dtype=np.float32) for _ in range(num_classes)]
                    for i, (bbox, score, label) in enumerate(zip(bboxes, scores, labels)):
                        if label < num_classes:
                            det = np.concatenate([bbox, np.array([score])], axis=0)
                            result_list[label] = np.vstack([result_list[label], det])

                    converted_results.append(result_list)

                results = converted_results

            # 计算每个类别的AP
            self.progress.emit("正在计算每个类别的AP...")
            class_ap_values = []
            class_recall_values = []

            try:
                mean_ap, ap_results = eval_map(results, annotations, iou_thr=0.5)

                for idx, cls_ap in enumerate(ap_results):
                    ap_value = cls_ap['ap']
                    recall_value = cls_ap['recall'][-1] if len(cls_ap['recall']) > 0 else 0
                    class_ap_values.append(ap_value)
                    class_recall_values.append(recall_value)
                    self.progress.emit(
                        f"类别 {class_names[idx]} (ID:{idx}) | AP@0.5 = {ap_value:.4f}, Recall = {recall_value:.4f}")

            except Exception as e:
                self.progress.emit(f"评估时出错: {e}")
                self.progress.emit("尝试手动计算每个类别的AP...")

                class_ap_values = []
                class_recall_values = []

                for class_id in range(num_classes):
                    self.progress.emit(f"计算类别 {class_names[class_id]} (ID:{class_id}) 的AP...")

                    class_results = []
                    class_annotations = []

                    for i, ann in enumerate(annotations):
                        mask = ann['labels'] == class_id
                        class_ann = {
                            'bboxes': ann['bboxes'][mask] if len(mask) > 0 else np.zeros((0, 4), dtype=np.float32),
                            'labels': np.zeros(np.sum(mask), dtype=np.int64)
                        }
                        class_annotations.append(class_ann)

                        if isinstance(results[i], list) and len(results[i]) > class_id:
                            class_results.append([results[i][class_id]])
                        else:
                            class_results.append([np.zeros((0, 5), dtype=np.float32)])

                    try:
                        class_mean_ap, class_ap_result = eval_map(
                            class_results,
                            class_annotations,
                            iou_thr=0.5
                        )
                        ap_value = class_mean_ap
                        recall_value = class_ap_result[0]['recall'][-1] if class_ap_result and len(
                            class_ap_result[0]['recall']) > 0 else 0

                        class_ap_values.append(ap_value)
                        class_recall_values.append(recall_value)

                        self.progress.emit(
                            f"类别 {class_names[class_id]} (ID:{class_id}) | AP@0.5 = {ap_value:.4f}, Recall = {recall_value:.4f}")
                    except Exception as e:
                        self.progress.emit(f"计算类别 {class_names[class_id]} 的AP时出错: {e}")
                        class_ap_values.append(0.0)
                        class_recall_values.append(0.0)

                mean_ap = sum(class_ap_values) / len(class_ap_values)

            # 返回结果
            self.progress.emit(f"评估完成! mAP@0.5 = {mean_ap:.4f}")

            results_dict = {
                'class_names': class_names,
                'class_ap_values': class_ap_values,
                'class_recall_values': class_recall_values,
                'mean_ap': mean_ap,
                'annotations': annotations,
                'results': results,
                'dataset': dataset
            }

            self.finished.emit(results_dict)

        except Exception as e:
            self.progress.emit(f"发生错误: {str(e)}")


class InferenceThread(QThread):
    finished = pyqtSignal(dict)
    progress = pyqtSignal(str)

    def __init__(self, model_path, config_path, image_path, confidence_threshold=0.5):
        super().__init__()
        self.model_path = model_path
        self.config_path = config_path
        self.image_path = image_path
        self.confidence_threshold = confidence_threshold

    def run(self):
        try:
            self.progress.emit("正在加载模型...")
            model = init_detector(self.config_path, self.model_path, device='cuda:0')

            self.progress.emit("正在进行推理...")
            result = inference_detector(model, self.image_path)

            self.progress.emit("推理完成!")

            # 获取类别名称
            class_names = model.dataset_meta['classes']

            # 处理结果
            pred_instances = result.pred_instances

            # 将结果转换为numpy数组
            bboxes = pred_instances.bboxes.cpu().numpy()
            scores = pred_instances.scores.cpu().numpy()
            labels = pred_instances.labels.cpu().numpy()

            # 过滤低置信度的检测结果
            mask = scores >= self.confidence_threshold
            bboxes = bboxes[mask]
            scores = scores[mask]
            labels = labels[mask]

            # 读取原始图像
            image = cv2.imread(self.image_path)
            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

            # 返回结果
            results_dict = {
                'image': image,
                'bboxes': bboxes,
                'scores': scores,
                'labels': labels,
                'class_names': class_names
            }

            self.finished.emit(results_dict)

        except Exception as e:
            self.progress.emit(f"推理时发生错误: {str(e)}")


class ModelEvaluationApp(QMainWindow):
    def __init__(self):
        super().__init__()

        self.setWindowTitle("目标检测模型评估工具")
        self.setGeometry(100, 100, 1200, 800)

        # 初始化变量
        self.model_path = ""
        self.config_path = ""
        self.result_path = ""
        self.image_dir = ""
        self.current_image_path = ""
        self.evaluation_results = None
        self.inference_results = None
        self.image_list = []
        self.current_image_index = 0

        # 创建主布局
        self.central_widget = QWidget()
        self.setCentralWidget(self.central_widget)

        self.main_layout = QVBoxLayout(self.central_widget)

        # 创建选项卡
        self.tabs = QTabWidget()
        self.main_layout.addWidget(self.tabs)

        # 创建评估选项卡
        self.evaluation_tab = QWidget()
        self.tabs.addTab(self.evaluation_tab, "模型评估")

        # 创建推理选项卡
        self.inference_tab = QWidget()
        self.tabs.addTab(self.inference_tab, "模型推理")

        # 设置评估选项卡
        self.setup_evaluation_tab()

        # 设置推理选项卡
        self.setup_inference_tab()

    def setup_evaluation_tab(self):
        layout = QVBoxLayout(self.evaluation_tab)

        # 创建文件选择区域
        file_group = QGroupBox("文件选择")
        file_layout = QGridLayout()

        # 模型路径
        self.model_label = QLabel("模型路径:")
        self.model_path_label = QLabel("未选择")
        self.model_browse_button = QPushButton("浏览...")
        self.model_browse_button.clicked.connect(self.browse_model)

        # 配置文件路径
        self.config_label = QLabel("配置文件:")
        self.config_path_label = QLabel("未选择")
        self.config_browse_button = QPushButton("浏览...")
        self.config_browse_button.clicked.connect(self.browse_config)

        # 结果文件路径
        self.result_label = QLabel("结果文件:")
        self.result_path_label = QLabel("未选择")
        self.result_browse_button = QPushButton("浏览...")
        self.result_browse_button.clicked.connect(self.browse_result)

        # 添加到布局
        file_layout.addWidget(self.model_label, 0, 0)
        file_layout.addWidget(self.model_path_label, 0, 1)
        file_layout.addWidget(self.model_browse_button, 0, 2)

        file_layout.addWidget(self.config_label, 1, 0)
        file_layout.addWidget(self.config_path_label, 1, 1)
        file_layout.addWidget(self.config_browse_button, 1, 2)

        file_layout.addWidget(self.result_label, 2, 0)
        file_layout.addWidget(self.result_path_label, 2, 1)
        file_layout.addWidget(self.result_browse_button, 2, 2)

        file_group.setLayout(file_layout)
        layout.addWidget(file_group)

        # 创建操作区域
        action_layout = QHBoxLayout()

        self.evaluate_button = QPushButton("开始评估")
        self.evaluate_button.clicked.connect(self.start_evaluation)
        action_layout.addWidget(self.evaluate_button)

        layout.addLayout(action_layout)

        # 创建进度区域
        progress_group = QGroupBox("评估进度")
        progress_layout = QVBoxLayout()

        self.progress_label = QLabel("等待开始评估...")
        progress_layout.addWidget(self.progress_label)

        progress_group.setLayout(progress_layout)
        layout.addWidget(progress_group)

        # 创建结果显示区域
        results_splitter = QSplitter(Qt.Horizontal)

        # 表格显示区域
        table_group = QGroupBox("类别准确率")
        table_layout = QVBoxLayout()

        self.results_table = QTableWidget()
        self.results_table.setColumnCount(3)
        self.results_table.setHorizontalHeaderLabels(["类别", "AP@0.5", "Recall@0.5"])
        self.results_table.setEditTriggers(QTableWidget.NoEditTriggers)

        table_layout.addWidget(self.results_table)
        table_group.setLayout(table_layout)
        results_splitter.addWidget(table_group)

        # 图表显示区域
        chart_group = QGroupBox("可视化结果")
        chart_layout = QVBoxLayout()

        self.chart_tabs = QTabWidget()

        # AP柱状图
        self.ap_chart_widget = QWidget()
        ap_chart_layout = QVBoxLayout(self.ap_chart_widget)
        self.ap_canvas = MatplotlibCanvas(width=5, height=4)
        ap_chart_layout.addWidget(self.ap_canvas)
        self.chart_tabs.addTab(self.ap_chart_widget, "AP柱状图")

        # AP和Recall对比图
        self.ap_recall_chart_widget = QWidget()
        ap_recall_chart_layout = QVBoxLayout(self.ap_recall_chart_widget)
        self.ap_recall_canvas = MatplotlibCanvas(width=5, height=4)
        ap_recall_chart_layout.addWidget(self.ap_recall_canvas)
        self.chart_tabs.addTab(self.ap_recall_chart_widget, "AP和Recall对比")

        # AP和Recall散点图
        self.scatter_chart_widget = QWidget()
        scatter_chart_layout = QVBoxLayout(self.scatter_chart_widget)
        self.scatter_canvas = MatplotlibCanvas(width=5, height=4)
        scatter_chart_layout.addWidget(self.scatter_canvas)
        self.chart_tabs.addTab(self.scatter_chart_widget, "AP和Recall散点图")

        chart_layout.addWidget(self.chart_tabs)
        chart_group.setLayout(chart_layout)
        results_splitter.addWidget(chart_group)

        layout.addWidget(results_splitter)

    def setup_inference_tab(self):
        layout = QVBoxLayout(self.inference_tab)

        # 创建文件选择区域
        file_group = QGroupBox("文件选择")
        file_layout = QGridLayout()

        # 模型路径
        self.inf_model_label = QLabel("模型路径:")
        self.inf_model_path_label = QLabel("未选择")
        self.inf_model_browse_button = QPushButton("浏览...")
        self.inf_model_browse_button.clicked.connect(self.browse_inf_model)

        # 配置文件路径
        self.inf_config_label = QLabel("配置文件:")
        self.inf_config_path_label = QLabel("未选择")
        self.inf_config_browse_button = QPushButton("浏览...")
        self.inf_config_browse_button.clicked.connect(self.browse_inf_config)

        # 图像文件夹路径
        self.image_dir_label = QLabel("图像文件夹:")
        self.image_dir_path_label = QLabel("未选择")
        self.image_dir_browse_button = QPushButton("浏览...")
        self.image_dir_browse_button.clicked.connect(self.browse_image_dir)

        # 添加到布局
        file_layout.addWidget(self.inf_model_label, 0, 0)
        file_layout.addWidget(self.inf_model_path_label, 0, 1)
        file_layout.addWidget(self.inf_model_browse_button, 0, 2)

        file_layout.addWidget(self.inf_config_label, 1, 0)
        file_layout.addWidget(self.inf_config_path_label, 1, 1)
        file_layout.addWidget(self.inf_config_browse_button, 1, 2)

        file_layout.addWidget(self.image_dir_label, 2, 0)
        file_layout.addWidget(self.image_dir_path_label, 2, 1)
        file_layout.addWidget(self.image_dir_browse_button, 2, 2)

        file_group.setLayout(file_layout)
        layout.addWidget(file_group)

        # 创建参数设置区域
        param_group = QGroupBox("参数设置")
        param_layout = QHBoxLayout()

        self.conf_threshold_label = QLabel("置信度阈值:")
        self.conf_threshold_slider = QSlider(Qt.Horizontal)
        self.conf_threshold_slider.setMinimum(0)
        self.conf_threshold_slider.setMaximum(100)
        self.conf_threshold_slider.setValue(50)
        self.conf_threshold_slider.setTickPosition(QSlider.TicksBelow)
        self.conf_threshold_slider.setTickInterval(10)

        self.conf_threshold_value = QLabel("0.50")
        self.conf_threshold_slider.valueChanged.connect(self.update_conf_threshold)

        param_layout.addWidget(self.conf_threshold_label)
        param_layout.addWidget(self.conf_threshold_slider)
        param_layout.addWidget(self.conf_threshold_value)

        param_group.setLayout(param_layout)
        layout.addWidget(param_group)

        # 创建操作区域
        action_layout = QHBoxLayout()

        self.prev_image_button = QPushButton("上一张")
        self.prev_image_button.clicked.connect(self.show_prev_image)
        self.prev_image_button.setEnabled(False)

        self.next_image_button = QPushButton("下一张")
        self.next_image_button.clicked.connect(self.show_next_image)
        self.next_image_button.setEnabled(False)

        self.infer_button = QPushButton("开始推理")
        self.infer_button.clicked.connect(self.start_inference)

        action_layout.addWidget(self.prev_image_button)
        action_layout.addWidget(self.infer_button)
        action_layout.addWidget(self.next_image_button)

        layout.addLayout(action_layout)

        # 创建进度区域
        inf_progress_group = QGroupBox("推理进度")
        inf_progress_layout = QVBoxLayout()

        self.inf_progress_label = QLabel("等待开始推理...")
        inf_progress_layout.addWidget(self.inf_progress_label)

        inf_progress_group.setLayout(inf_progress_layout)
        layout.addWidget(inf_progress_group)

        # 创建图像显示区域 - 使用水平分割器
        image_splitter = QSplitter(Qt.Horizontal)

        # 左侧：原始图像
        original_group = QGroupBox("原始图像")
        original_layout = QVBoxLayout()

        self.original_scroll = QScrollArea()
        self.original_scroll.setWidgetResizable(True)
        self.original_scroll.setAlignment(Qt.AlignCenter)

        self.original_label = QLabel()
        self.original_label.setAlignment(Qt.AlignCenter)
        self.original_label.setMinimumSize(400, 300)

        self.original_scroll.setWidget(self.original_label)
        original_layout.addWidget(self.original_scroll)
        original_group.setLayout(original_layout)

        # 右侧：检测结果
        result_group = QGroupBox("检测结果")
        result_layout = QVBoxLayout()

        self.result_scroll = QScrollArea()
        self.result_scroll.setWidgetResizable(True)
        self.result_scroll.setAlignment(Qt.AlignCenter)

        self.result_label = QLabel()
        self.result_label.setAlignment(Qt.AlignCenter)
        self.result_label.setMinimumSize(400, 300)

        self.result_scroll.setWidget(self.result_label)
        result_layout.addWidget(self.result_scroll)
        result_group.setLayout(result_layout)

        # 添加到分割器
        image_splitter.addWidget(original_group)
        image_splitter.addWidget(result_group)

        # 设置初始分割比例
        image_splitter.setSizes([600, 600])

        layout.addWidget(image_splitter)

    def browse_model(self):
        file_path, _ = QFileDialog.getOpenFileName(self, "选择模型文件", "", "PyTorch模型 (*.pth)")
        if file_path:
            self.model_path = file_path
            self.model_path_label.setText(os.path.basename(file_path))

    def browse_config(self):
        file_path, _ = QFileDialog.getOpenFileName(self, "选择配置文件", "", "Python文件 (*.py)")
        if file_path:
            self.config_path = file_path
            self.config_path_label.setText(os.path.basename(file_path))

    def browse_result(self):
        file_path, _ = QFileDialog.getOpenFileName(self, "选择结果文件", "", "Pickle文件 (*.pkl)")
        if file_path:
            self.result_path = file_path
            self.result_path_label.setText(os.path.basename(file_path))

    def browse_inf_model(self):
        file_path, _ = QFileDialog.getOpenFileName(self, "选择模型文件", "", "PyTorch模型 (*.pth)")
        if file_path:
            self.model_path = file_path
            self.inf_model_path_label.setText(os.path.basename(file_path))

    def browse_inf_config(self):
        file_path, _ = QFileDialog.getOpenFileName(self, "选择配置文件", "", "Python文件 (*.py)")
        if file_path:
            self.config_path = file_path
            self.inf_config_path_label.setText(os.path.basename(file_path))

    def browse_image_dir(self):
        dir_path = QFileDialog.getExistingDirectory(self, "选择图像文件夹")
        if dir_path:
            self.image_dir = dir_path
            self.image_dir_path_label.setText(os.path.basename(dir_path))

            # 获取图像文件列表
            self.image_list = []
            for file in os.listdir(dir_path):
                if file.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.tif', '.tiff')):
                    self.image_list.append(os.path.join(dir_path, file))

            if self.image_list:
                self.current_image_index = 0
                self.current_image_path = self.image_list[0]
                self.prev_image_button.setEnabled(False)
                self.next_image_button.setEnabled(len(self.image_list) > 1)

                # 显示第一张图像
                self.display_image(self.current_image_path)

    def update_conf_threshold(self):
        value = self.conf_threshold_slider.value() / 100.0
        self.conf_threshold_value.setText(f"{value:.2f}")

        # 如果已经有推理结果，重新显示图像
        if self.inference_results is not None and self.current_image_path:
            self.display_inference_results(self.inference_results, value)

    def start_evaluation(self):
        if not self.result_path or not self.config_path:
            self.progress_label.setText("错误: 请选择结果文件和配置文件")
            return

        self.progress_label.setText("正在开始评估...")
        self.evaluate_button.setEnabled(False)

        # 创建并启动评估线程
        self.eval_thread = EvaluationThread(self.result_path, self.config_path)
        self.eval_thread.progress.connect(self.update_eval_progress)
        self.eval_thread.finished.connect(self.evaluation_finished)
        self.eval_thread.start()

    def update_eval_progress(self, message):
        self.progress_label.setText(message)

    def evaluation_finished(self, results):
        self.evaluation_results = results
        self.evaluate_button.setEnabled(True)

        # 更新表格
        class_names = results['class_names']
        class_ap_values = results['class_ap_values']
        class_recall_values = results['class_recall_values']

        self.results_table.setRowCount(len(class_names))

        for i, (name, ap, recall) in enumerate(zip(class_names, class_ap_values, class_recall_values)):
            self.results_table.setItem(i, 0, QTableWidgetItem(name))
            self.results_table.setItem(i, 1, QTableWidgetItem(f"{ap:.4f}"))
            self.results_table.setItem(i, 2, QTableWidgetItem(f"{recall:.4f}"))

        # 添加mAP行
        self.results_table.setRowCount(len(class_names) + 1)
        self.results_table.setItem(len(class_names), 0, QTableWidgetItem("mAP"))
        self.results_table.setItem(len(class_names), 1, QTableWidgetItem(f"{results['mean_ap']:.4f}"))
        self.results_table.setItem(len(class_names), 2, QTableWidgetItem("-"))

        # 调整表格列宽
        self.results_table.resizeColumnsToContents()

        # 更新图表
        self.update_charts(results)

    def update_charts(self, results):
        class_names = results['class_names']
        class_ap_values = results['class_ap_values']
        class_recall_values = results['class_recall_values']
        mean_ap = results['mean_ap']

        # 更新AP柱状图
        self.ap_canvas.axes.clear()
        bars = self.ap_canvas.axes.bar(class_names, class_ap_values, color='skyblue')
        self.ap_canvas.axes.axhline(y=mean_ap, color='r', linestyle='-', label=f'mAP = {mean_ap:.4f}')
        self.ap_canvas.axes.set_xlabel('类别')
        self.ap_canvas.axes.set_ylabel('AP@0.5')
        self.ap_canvas.axes.set_title('每个类别的AP@0.5')
        self.ap_canvas.axes.set_xticklabels(class_names, rotation=45, ha='right')
        self.ap_canvas.axes.set_ylim(0, 1.0)
        self.ap_canvas.axes.grid(axis='y', linestyle='--', alpha=0.7)

        # 在每个柱子上方添加AP值
        for bar, ap_value in zip(bars, class_ap_values):
            self.ap_canvas.axes.text(bar.get_x() + bar.get_width() / 2, bar.get_height() + 0.02,
                                     f'{ap_value:.3f}', ha='center', fontsize=9)

        self.ap_canvas.axes.legend()
        self.ap_canvas.fig.tight_layout()
        # 更新AP和Recall对比图
        self.ap_recall_canvas.axes.clear()
        x = np.arange(len(class_names))
        width = 0.35

        bars1 = self.ap_recall_canvas.axes.bar(x - width / 2, class_ap_values, width, label='AP@0.5', color='skyblue')
        bars2 = self.ap_recall_canvas.axes.bar(x + width / 2, class_recall_values, width, label='Recall@0.5',
                                               color='lightgreen')

        self.ap_recall_canvas.axes.set_xlabel('类别')
        self.ap_recall_canvas.axes.set_ylabel('分数')
        self.ap_recall_canvas.axes.set_title('每个类别的AP@0.5和Recall@0.5对比')
        self.ap_recall_canvas.axes.set_xticks(x)
        self.ap_recall_canvas.axes.set_xticklabels(class_names, rotation=45, ha='right')
        self.ap_recall_canvas.axes.set_ylim(0, 1.0)
        self.ap_recall_canvas.axes.grid(axis='y', linestyle='--', alpha=0.7)
        self.ap_recall_canvas.axes.legend()

        # 在每个柱子上方添加值
        for bar, value in zip(bars1, class_ap_values):
            self.ap_recall_canvas.axes.text(bar.get_x() + bar.get_width() / 2, bar.get_height() + 0.02,
                                            f'{value:.2f}', ha='center', fontsize=8)
        for bar, value in zip(bars2, class_recall_values):
            self.ap_recall_canvas.axes.text(bar.get_x() + bar.get_width() / 2, bar.get_height() + 0.02,
                                            f'{value:.2f}', ha='center', fontsize=8)

        self.ap_recall_canvas.fig.tight_layout()

        # 更新AP和Recall散点图
        self.scatter_canvas.axes.clear()
        self.scatter_canvas.axes.scatter(class_ap_values, class_recall_values, color='blue', s=100, alpha=0.7)

        # 添加类别标签
        for i, txt in enumerate(class_names):
            self.scatter_canvas.axes.annotate(txt, (class_ap_values[i], class_recall_values[i]),
                                              xytext=(5, 5), textcoords='offset points', fontsize=10)

        self.scatter_canvas.axes.set_xlabel('AP@0.5')
        self.scatter_canvas.axes.set_ylabel('Recall@0.5')
        self.scatter_canvas.axes.set_title('AP@0.5 vs Recall@0.5散点图')
        self.scatter_canvas.axes.grid(True, linestyle='--', alpha=0.7)
        self.scatter_canvas.axes.set_xlim(0, 1.0)
        self.scatter_canvas.axes.set_ylim(0, 1.0)

        # 添加对角线参考
        self.scatter_canvas.axes.plot([0, 1], [0, 1], 'r--', alpha=0.5)

        self.scatter_canvas.fig.tight_layout()

        # 刷新画布
        self.ap_canvas.draw()
        self.ap_recall_canvas.draw()
        self.scatter_canvas.draw()

    def start_inference(self):
        if not self.model_path or not self.config_path or not self.current_image_path:
            self.inf_progress_label.setText("错误: 请选择模型、配置文件和图像")
            return

        self.inf_progress_label.setText("正在开始推理...")
        self.infer_button.setEnabled(False)

        # 获取当前置信度阈值
        confidence_threshold = self.conf_threshold_slider.value() / 100.0

        # 创建并启动推理线程
        self.inf_thread = InferenceThread(
            self.model_path,
            self.config_path,
            self.current_image_path,
            confidence_threshold
        )
        self.inf_thread.progress.connect(self.update_inf_progress)
        self.inf_thread.finished.connect(self.inference_finished)
        self.inf_thread.start()

    def update_inf_progress(self, message):
        self.inf_progress_label.setText(message)

    def inference_finished(self, results):
        self.inference_results = results
        self.infer_button.setEnabled(True)

        # 显示推理结果
        confidence_threshold = self.conf_threshold_slider.value() / 100.0
        self.display_inference_results(results, confidence_threshold)

    def display_inference_results(self, results, confidence_threshold=0.5):
        # 显示原始图像
        original_pixmap = self.create_pixmap_from_array(results['image'])
        self.display_pixmap(self.original_label, self.original_scroll, original_pixmap)

        # 创建检测结果图像
        image = results['image'].copy()
        bboxes = results['bboxes']
        scores = results['scores']
        labels = results['labels']
        class_names = results['class_names']

        # 绘制检测框
        height, width, _ = image.shape
        font = cv2.FONT_HERSHEY_SIMPLEX
        font_scale = 0.5
        thickness = 2

        # 定义颜色映射
        colors = [
            (255, 0, 0),  # 红色
            (0, 255, 0),  # 绿色
            (0, 0, 255),  # 蓝色
            (255, 255, 0),  # 黄色
            (255, 0, 255),  # 品红
            (0, 255, 255),  # 青色
            (128, 0, 0),  # 深红
            (0, 128, 0),  # 深绿
            (0, 0, 128),  # 深蓝
            (128, 128, 0)  # 橄榄
        ]

        # 过滤低置信度的检测结果
        mask = scores >= confidence_threshold
        bboxes = bboxes[mask]
        scores = scores[mask]
        labels = labels[mask]

        for bbox, score, label in zip(bboxes, scores, labels):
            x1, y1, x2, y2 = map(int, bbox)

            # 确保边界框在图像范围内
            x1 = max(0, x1)
            y1 = max(0, y1)
            x2 = min(width - 1, x2)
            y2 = min(height - 1, y2)

            # 获取类别颜色
            color = colors[int(label) % len(colors)]

            # 绘制边界框
            cv2.rectangle(image, (x1, y1), (x2, y2), color, thickness)

            # 准备标签文本
            if int(label) < len(class_names):
                label_text = f"{class_names[int(label)]}: {score:.2f}"
            else:
                label_text = f"Unknown: {score:.2f}"

            # 计算文本大小
            (text_width, text_height), _ = cv2.getTextSize(label_text, font, font_scale, thickness)

            # 绘制文本背景
            cv2.rectangle(image, (x1, y1 - text_height - 5), (x1 + text_width, y1), color, -1)

            # 绘制文本
            cv2.putText(image, label_text, (x1, y1 - 5), font, font_scale, (255, 255, 255), thickness)

        # 显示检测结果图像
        result_pixmap = self.create_pixmap_from_array(image)
        self.display_pixmap(self.result_label, self.result_scroll, result_pixmap)

    def create_pixmap_from_array(self, image):
        """从numpy数组创建QPixmap"""
        height, width, channel = image.shape
        bytes_per_line = 3 * width
        q_image = QImage(image.data, width, height, bytes_per_line, QImage.Format_RGB888)
        return QPixmap.fromImage(q_image)

    def display_pixmap(self, label, scroll_area, pixmap):
        """自适应显示QPixmap到指定的标签和滚动区域"""
        # 获取滚动区域的大小
        scroll_size = scroll_area.viewport().size()

        # 计算缩放比例，确保图像适应滚动区域
        scale_factor = min(
            scroll_size.width() / pixmap.width(),
            scroll_size.height() / pixmap.height()
        )

        # 如果图像比滚动区域小，则不缩放
        if scale_factor > 1:
            scale_factor = 1

        # 缩放图像
        if scale_factor < 1:
            scaled_pixmap = pixmap.scaled(
                int(pixmap.width() * scale_factor),
                int(pixmap.height() * scale_factor),
                Qt.KeepAspectRatio,
                Qt.SmoothTransformation
            )
        else:
            scaled_pixmap = pixmap

        # 显示图像
        label.setPixmap(scaled_pixmap)
        label.setMinimumSize(scaled_pixmap.width(), scaled_pixmap.height())

        # 更新状态栏显示图像信息
        self.statusBar().showMessage(
            f"图像大小: {pixmap.width()}x{pixmap.height()} | 显示大小: {scaled_pixmap.width()}x{scaled_pixmap.height()}")

    def display_image(self, image_path):
        if not os.path.exists(image_path):
            return

        # 读取原始图像
        image = cv2.imread(image_path)
        if image is None:
            return

        # 转换为RGB
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

        # 显示原始图像
        original_pixmap = self.create_pixmap_from_array(image)
        self.display_pixmap(self.original_label, self.original_scroll, original_pixmap)

        # 清除结果图像
        self.result_label.clear()
        self.result_label.setText("等待推理结果...")

    def show_prev_image(self):
        if not self.image_list or self.current_image_index <= 0:
            return

        self.current_image_index -= 1
        self.current_image_path = self.image_list[self.current_image_index]

        # 更新按钮状态
        self.prev_image_button.setEnabled(self.current_image_index > 0)
        self.next_image_button.setEnabled(self.current_image_index < len(self.image_list) - 1)

        # 显示图像
        self.display_image(self.current_image_path)

        # 清除之前的推理结果
        self.inference_results = None

    def show_next_image(self):
        if not self.image_list or self.current_image_index >= len(self.image_list) - 1:
            return

        self.current_image_index += 1
        self.current_image_path = self.image_list[self.current_image_index]

        # 更新按钮状态
        self.prev_image_button.setEnabled(self.current_image_index > 0)
        self.next_image_button.setEnabled(self.current_image_index < len(self.image_list) - 1)

        # 显示图像
        self.display_image(self.current_image_path)

        # 清除之前的推理结果
        self.inference_results = None

    def resizeEvent(self, event):
        """窗口大小改变时自适应调整图像大小"""
        super().resizeEvent(event)

        # 如果有图像正在显示，则重新调整大小
        if hasattr(self,
                   'original_label') and self.original_label.pixmap() and not self.original_label.pixmap().isNull():
            self.display_pixmap(self.original_label, self.original_scroll, self.original_label.pixmap().copy())

        if hasattr(self, 'result_label') and self.result_label.pixmap() and not self.result_label.pixmap().isNull():
            self.display_pixmap(self.result_label, self.result_scroll, self.result_label.pixmap().copy())


if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = ModelEvaluationApp()
    window.show()
    sys.exit(app.exec_())

